from tensorflow.keras import Input, Model
from tensorflow.keras.layers import Lambda, Conv2D, Activation, MaxPooling2D, Flatten, Dense, Dropout, \
    BatchNormalization
import tensorflow as tf

from src.image_classification import data_processor


def network_cnn(input_shape, num_classes):
    """
    卷积神经网络

    :param input_shape: 输入的图片shape
    :param num_classes: 最大的分类数目
    :return: model
    """
    img_input = Input(shape=input_shape, name='data')
    x = Lambda(data_processor.convert_to_hsv_and_grayscale)(img_input)
    x = Conv2D(16, (5, 5), strides=(1, 1), padding='same', name='conv1')(x)
    x = Activation('relu', name='conv1_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool1')(x)
    x = Conv2D(32, (5, 5), strides=(1, 1), padding='same', name='conv2')(x)
    x = Activation('relu', name='conv2_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool2')(x)
    x = Conv2D(64, (5, 5), strides=(1, 1), padding='same', name='conv3')(x)
    x = Activation('relu', name='conv3_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool3')(x)
    x = Conv2D(128, (5, 5), strides=(1, 1), padding='same', name='conv4')(x)
    x = Activation('relu', name='conv4_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool4')(x)
    x = Flatten()(x)
    x = Dense(1024, activation='relu', name='fcl1')(x)
    x = Dropout(0.2)(x)
    x = Dense(256, activation='relu', name='fcl2')(x)
    x = Dropout(0.2)(x)
    out = Dense(num_classes, activation='softmax', name='predictions')(x)
    rez = Model(inputs=img_input, outputs=out)
    return rez


def network_cnn_normalization(input_shape, num_classes):
    """
    卷积神经网络(进行了归一化的)

    :param input_shape: 输入的图片shape
    :param num_classes: 最大的分类数目
    :return: model
    """
    img_input = Input(shape=input_shape, name='data')
    x = Lambda(data_processor.convert_to_hsv_and_grayscale)(img_input)
    x = Conv2D(16, (5, 5), strides=(1, 1), padding='same', name='conv1')(x)
    x = Activation('relu', name='conv1_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool1')(x)
    x = Conv2D(32, (5, 5), strides=(1, 1), padding='same', name='conv2')(x)
    x = Activation('relu', name='conv2_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool2')(x)
    x = Conv2D(64, (5, 5), strides=(1, 1), padding='same', name='conv3')(x)
    x = Activation('relu', name='conv3_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool3')(x)
    x = Conv2D(128, (5, 5), strides=(1, 1), padding='same', name='conv4')(x)
    x = Activation('relu', name='conv4_relu')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), padding='valid', name='pool4')(x)
    x = Flatten()(x)
    x = BatchNormalization()(x)
    x = Dense(1024, name='fcl1')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Dropout(0.2)(x)
    x = Dense(256, name='fcl2')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Dropout(0.2)(x)
    out = Dense(num_classes, activation='softmax', name='predictions')(x)
    rez = Model(inputs=img_input, outputs=out)
    return rez


def network_resnet(input_shape, num_classes):
    """
    resnet

    :param input_shape: 输入的图片shape
    :param num_classes: 最大的分类数目
    :return: model
    """
    img_input = Input(shape=input_shape, name='data')
    x = Lambda(data_processor.convert_to_hsv_and_grayscale)(img_input)
    model = tf.keras.applications.ResNet152(include_top=True, input_tensor=x, pooling='max',
                                            classes=num_classes, weights=None)
    return model


if __name__ == '__main__':
    print(network_resnet((100, 100, 3), 100).summary())
